/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2017-2021. All rights reserved.
 * Description: The description of the MemoryAllocatorFactory in framework
 */
#ifndef _DOMI_COMMMON_MEMORY_ALLOCATOR_FACTORY_H_
#define _DOMI_COMMMON_MEMORY_ALLOCATOR_FACTORY_H_

#include <string>
#include <memory>
#include <mutex>
#include <map>
#include <functional>

#include "framework/common/memory_allocator.h"
#include "framework/common/hcs_types.h"

namespace hiai {
using CREATOR_FUN = std::function<std::shared_ptr<hiai::MemoryAllocator>()>;

// MemoryAllocatorFactory: manage all the allocator, support create allocator from MemoryAllocator
class HCS_API_EXPORT MemoryAllocatorFactory {
public:
    // this is a singleton object
    static MemoryAllocatorFactory* Instance();

    /* @ingroup hiai
     * @brief 创建allocator
     * @param [in] type 内存类型
     * @return 非nullptr 成功
     * @return nullptr 失败
     */
    std::shared_ptr<hiai::MemoryAllocator> CreateAllocator(hiai::MemoryType type);
    /* @ingroup domi_ome
     * @brief 注册Op创建函数，这个接口在MemoryAllocator的实现类中通过REGISTER_ALLOCATOR_CREATOR调用
     * @param [in] type allocator类型
     * @param [in] fun allocator创建函数
     */
    void RegisterAllocator(hiai::MemoryType type, CREATOR_FUN fun);

protected:
    MemoryAllocatorFactory() = default;
    ~MemoryAllocatorFactory() = default;

private:
    // the allocator creator funtion map
    std::map<hiai::MemoryType, CREATOR_FUN> allocatorCreatorMap_;
}; // end class MemoryAllocatorFactory

class AllocatorRegisterar {
public:
    AllocatorRegisterar(hiai::MemoryType type, CREATOR_FUN fun)
    {
        MemoryAllocatorFactory::Instance()->RegisterAllocator(type, fun);
    }
    ~AllocatorRegisterar()
    {
    }
};

#define REGISTER_ALLOCATOR_CREATOR(type, zoneType, clazz) \
    std::shared_ptr<hiai::MemoryAllocator> Creator_##type##_Allocator() \
    { \
        std::shared_ptr<clazz> ptr = std::shared_ptr<clazz>(new (std::nothrow) clazz()); \
        return std::shared_ptr<hiai::MemoryAllocator>(ptr); \
    } \
    const AllocatorRegisterar g_##type##_allocator_creator(zoneType, Creator_##type##_Allocator)
}; // end namespace hiai
#endif
